import cv2
import os
import numpy as np
from torch.utils.data import Dataset
from torchvision import transforms
import glob
import tqdm
import pickle

mean = [0.449, 0.436, 0.421]
std = [0.284, 0.287, 0.306]

class ParisRetriData(Dataset):
    def __init__(self, paris_path):
        super().__init__()
        self.images = glob.glob(os.path.join(paris_path, "jpg/*.jpg"))
        # gt is a list ['gnd', 'imlist', 'qimlist']
        self.gt =  pickle.load(open(os.path.join(paris_path, 'gnd_rparis6k.pkl'), 'rb'))
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)
        ])
        self.query_list = self.create_query_list()  # 没有jpg后缀的文件名
        self.query_map = dict()
        self.create_query_maps()
    
    # 读取query_list数据
    """
        注意，'gnd'与'qimlist'的长度是一致且一一对应的
    """
    def create_query_list(self):
        query_list = [f"{file}.jpg" for file in self.gt['qimlist']]
        return query_list

    def create_query_maps(self):
        """
        This creates a dictionary of dictionary that contains the postive and negative examples for every query
        """
        # query是文件名称，如 all_souls_000013
        for i in range(len(self.query_list)):
            # 创建临时字典
            print(self.query_list[i].split('.')[0])
            tmp = dict()
            # 取对应 self.gt['imlist'] 中的索引
            easy_file_idx, hard_file_idx, junk_file_idx = self.gt['gnd'][i]['easy'], self.gt['gnd'][i]['hard'], self.gt['gnd'][i]['junk']
            # 都是索引值，是个数字，需要自己取 self.gt['imlist'] 相应的图片
            easy_file = [f"{self.gt['imlist'][i]}.jpg" for i in easy_file_idx]
            hard_file = [f"{self.gt['imlist'][i]}.jpg" for i in hard_file_idx]
            junk_file = [f"{self.gt['imlist'][i]}.jpg" for i in junk_file_idx]
            tmp['positive'] =  easy_file + hard_file + junk_file
            print("> Creating hard negative samples for query:", self.query_list[i].split('.')[0])
            tmp['negative'] = self._get_remaining_image_files(set(tmp['positive'] + [self.query_list[i]] + self._get_blacklist()))
            tmp['negative'] = self._create_bad_image_files(query, query_image_name, tmp['negative'])
    
    def _get_all_image_files(self):
        all_file_list = [file for file in self.images]
        return all_file_list

    def _get_remaining_image_files(self, tmp_set):
        """
        Get all the negative images corresponding to query
        """
        all_set = set(self._get_all_image_files())
        bad_list = list(all_set - tmp_set)
        return bad_list
    
    def _get_blacklist(self):
        """
        Paris 6k dataset has blacklisted images that should be filtered.
        """
        return ["paris_louvre_000136.jpg",
        "paris_louvre_000146.jpg",
        "paris_moulinrouge_000422.jpg",
        "paris_museedorsay_001059.jpg",
        "paris_notredame_000188.jpg",
        "paris_pantheon_000284.jpg",
        "paris_pantheon_000960.jpg",
        "paris_pantheon_000974.jpg",
        "paris_pompidou_000195.jpg",
        "paris_pompidou_000196.jpg",
        "paris_pompidou_000201.jpg",
        "paris_pompidou_000467.jpg",
        "paris_pompidou_000640.jpg",
        "paris_sacrecoeur_000299.jpg",
        "paris_sacrecoeur_000330.jpg",
        "paris_sacrecoeur_000353.jpg",
        "paris_triomphe_000662.jpg",
        "paris_triomphe_000833.jpg",
        "paris_triomphe_000863.jpg",
        "paris_triomphe_000867.jpg",]

    # def _create_bad_image_files(self, )

    def __getitem__(self, i):
        print(self.images[i])
        img = cv2.imread(self.images[i])
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = self.transform(img)
        return img

    def __len__(self):
        return len(self.images)

def loadparis(path):
    return ParisRetriData(path)

if __name__ == "__main__":
    
    # 删除无效数据
    # data = glob.glob('/workspace/wzj/revisitop/data/datasets/rparis6k/jpg/*.jpg')
    # for i in tqdm.tqdm(data):
    #     img = cv2.imread(i)
    #     if type(img) == np.ndarray:
    #         continue
    #     os.remove(i)
    #     print(f"remove {i}")
            

    # img = cv2.imread('/workspace/wzj/revisitop/data/datasets/rparis6k/jpg/paris_pompidou_000640.jpg')
    # img = cv2.imread('/workspace/wzj/revisitop/data/datasets/rparis6k/jpg/paris_defense_000000.jpg')
    # print(type(img))
    data = ParisRetriData('/workspace/wzj/revisitop/data/datasets/rparis6k')
    print(data.query_map)